package com.ds.lens.data.common.query.es;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLLimit;
import com.alibaba.druid.sql.ast.SQLObject;
import com.alibaba.druid.sql.ast.expr.*;
import com.alibaba.druid.sql.ast.statement.SQLSelectGroupByClause;
import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock;
import com.alibaba.druid.sql.visitor.SQLASTOutputVisitor;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class MQL2ElasticSQLAdaptor extends SQLASTOutputVisitor {
    private static final Map<String, TranslateFunction> FUNCTION_MAP;
    private Number limit;
    private SQLBinaryOpExpr where;

    static {
        FUNCTION_MAP = new HashMap<>();
        FUNCTION_MAP.put(DateHistogramFunction.NAME, new DateHistogramFunction());
        FUNCTION_MAP.put(TermsFunction.NAME, new TermsFunction());
        FUNCTION_MAP.put(TermFunction.NAME, new TermFunction());
        FUNCTION_MAP.put(InTermsFunction.NAME, new InTermsFunction());
        FUNCTION_MAP.put(NoFunction.NAME, new NoFunction());
    }

    public MQL2ElasticSQLAdaptor(Appendable appender) {
        super(appender);
    }

    public boolean visit(SQLSelectQueryBlock x) {
        SQLLimit sqlLimit = x.getLimit();
        if (sqlLimit != null) {
            limit = ((SQLIntegerExpr) sqlLimit.getRowCount()).getNumber();
        }
        where = (SQLBinaryOpExpr) x.getWhere();
        return super.visit(x);
    }

    public boolean visit(SQLMethodInvokeExpr x) {
        TranslateFunction function = FUNCTION_MAP.get(x.getMethodName());
        if (function == null) {
            return super.visit(x);
        }
        try {
            List<SQLExpr> exprList = x.getParameters();
            if (exprList != null && !exprList.isEmpty()) {
                String[] paramArr = new String[exprList.size()];
                for (int i = 0; i < exprList.size(); i++) {
                    paramArr[i] = exprList.get(i).toString();
                }
                function.apply(this.appender, paramArr);
            } else {
                function.apply(this.appender, x.getMethodName());
            }
            return false;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public boolean visit(SQLIdentifierExpr x) {
        if (limit != null && x.getParent() instanceof SQLSelectGroupByClause) {
            try {
                TranslateFunction function = FUNCTION_MAP.get(TermsFunction.NAME);
                function.apply(appender, x.getName(), limit.toString());
                return false;
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        } else {
            return super.visit(x);
        }
    }

    public boolean visit(SQLInListExpr x) {
        if (!x.isNot() && isInWhereClause(x)) {
            try {
                List<SQLExpr> exprList = x.getTargetList();
                String[] paramArr = new String[exprList.size() + 1];
                paramArr[0] = x.getExpr().toString();
                for (int i = 0; i < exprList.size(); i++) {
                    paramArr[i + 1] = exprList.get(i).toString();
                }
                TranslateFunction function = FUNCTION_MAP.get(InTermsFunction.NAME);
                function.apply(appender, paramArr);
                return false;
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
        return super.visit(x);
    }

    public boolean visit(SQLBinaryOpExpr x) {
        if (SQLBinaryOperator.Equality.equals(x.getOperator()) && isInWhereClause(x)) {
            try {
                if ("1".equals(x.getLeft().toString()) && "1".equals(x.getRight().toString())) {
                    TranslateFunction function = FUNCTION_MAP.get(NoFunction.NAME);
                    function.apply(appender, x.getLeft().toString(), x.getRight().toString());
                    return false;
                }
                TranslateFunction function = FUNCTION_MAP.get(TermFunction.NAME);
                function.apply(appender, x.getLeft().toString(), x.getRight().toString());
                return false;
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
        return super.visit(x);
    }

    private boolean isInWhereClause(SQLObject x) {
        SQLObject parent = x;
        while (where != null && !where.equals(parent) && parent != null) {
            parent = parent.getParent();
        }
        return parent != null;
    }

}
